#!/usr/bin/env python 

import os
import sys
import random
import tempfile

###
# Try to import things that don't exist on Windows

try:
    import grp
except ImportError:
    grp = None

try:
    import pwd
except ImportError:
    pwd = None

try:
    import stat
except ImportError:
    stat = None

try:
    import fcntl
except ImportError:
    fcntl = None
    pass

"""pyfileutils is a collection of safe, conservative file utilities.

It is inconvenient to worry about flushing/closing/cleaning up when
handling files, so modules that have to work with files can simply
import the wrapper functions contained herein and clean up the code a
little.  Also, operations like copy, move, getting basenames/dirnames
of files, listing the full paths of a folder's contents, etc. are
over-complicated or over-simple.  Or, worse, if you want different
behavior you have to re-implement the code with your change.

pyfileutils will raise OverwriteError if files are about to be
truncated or overwritten, but this behavior is overridable.  On POSIX
systems, file locking is available.

Exceptions are still thrown as normal, the usual ones being OSError and
IOError.

When creating files or folders, the following methods are used:

    example_func(name=None, dir=None)

- If name is omitted, it is randomly generated using the SHA-1 hash of
  a random string.
- If dir is omitted
  - If name was omitted, the system temporary folder is used as a
    parent folder.
  - Otherwise the current working directory is used.
- If dir is a relative path, the current working directory is used as
  the base folder.
- If dir is an absolute path, it is used as given.

"""

def get_random_string(method='sha-1'):
    """Returns a string containing random characters.

    method: a string representing which method to use.  Supported
            methods include:
              - sha-1:   returns an SHA-1 hash, of a "random" string;
                         the hash is 40 characters long.
              - mkstemp: returns what mkstemp would return, a "random"
                         string 6 characters long.
    
    """
    x = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_'
    if method == 'sha-1':
        try:
            import hashlib
            hash = hashlib.sha1()
        except ImportError:
            import sha
            hash = sha.new()
        r = range(52)
        [hash.update(''.join([random.choice(x) for y in r])) for z in r]
        return hash.hexdigest()
    elif method == 'mkstemp':
        return ''.join([random.choice(x) for y in range(6)])
    else:
        raise ValueError("Unsupported method [%s]" % (method))

###
#
# TODO:
#
#     - Add basename/dirname implementations that work.
#
#     - Add copy and move commands that work across filesystems, if not
#           atomically (remove the new file and raise an Exception if the
#           files differ and operation was not atomic).
#
###

class OverwriteError(Exception):
    def __init__(self, existing):
        Exception.__init__(self, "Refusing to overwrite [%s]" % (existing))

__all__ = ['ls', 'create_file', 'read_file', 'write_file', 'update_file',
           'append_file', 'delete_file', 'create_folder', 'delete_folder',
           'rmtree', 'mktree', 'get_random_string', 'get_random_file',
           'get_temp_file', 'get_temp_dir', 'OverwriteError']

def _get_path(name=None, dir=None, prefix=None, suffix=None, method='sha-1'):
    if name is None:
        if dir is None:
            dir == tempfile.gettempdir()
        name = get_random_string(method=method)
    if os.path.isabs(name):
        if dir is not None:
            es = "Received both an absolute file path and a target folder"
            raise ValueError(es)
        else:
            dir = os.path.dirname(name)
    else:
        dir = dir or os.getcwd()
    if prefix is not None:
        name = prefix + name
    if suffix is not None:
        name += suffix
    return os.path.join(dir, name)

from File import File
from LockFile import LockFile

def chmod(x, mode):
    """Changes the permissions on a file/folder.

    x:    a string representing the full path to a file or folder.
    mode: a octal number representing the new mode to set.

    """
    mode = oct(mode) # turns it into a string, strangely
    if len(mode) == 4:
        mode = mode[1:]
    user, group, other = mode
    new_mode = 0
    if user & 4:
        new_mode |= stat.S_IRUSR
    if user & 2:
        new_mode |= stat.S_IWUSR
    if user & 1:
        new_mode |= stat.S_IXUSR
    if group & 4:
        new_mode |= stat.S_IRGRP
    if group & 2:
        new_mode |= stat.S_IWGRP
    if group & 1:
        new_mode |= stat.S_IXGRP
    if other & 4:
        new_mode |= stat.S_IROTH
    if other & 2:
        new_mode |= stat.S_IWOTH
    if other & 1:
        new_mode |= stat.S_IXOTH
    os.chmod(x, new_mode)

def chown(x, uid=None, gid=None):
    """Sets the owner/group of a file/folder.

    x:   a string representing the full path to a file or folder.
    uid: a string representing the name of the new user.
         Alternately, an int representing the uid of the new user.
    gid: a string representing the name of the new group.
         Alternately, an int representing the gid of the new group.

    UID and GID are optional, but giving neither is a no-op.

    """
    if uid is None and gid is None:
        return
    s = os.stat(x)
    uid = uid or s.st_uid
    gid = gid or s.st_gid
    if isinstance(uid, str):
        uid = pwd.getpwnam(owner).pw_uid
    if isinstance(gid, str):
        gid = grp.getgrnam(group).gr_gid
    os.chown(x, uid, gid)

def chgrp(x, group):
    """Changes the group of a file/folder.

    x:     a string representing the full path to a file or folder.
    group: a string representing the name of the new group.
           Alternately, an int representing the gid of the new group.

    """
    chown(x, None, group)

def chmod_tree(x, mode):
    """Sets permissions on a folder and all its contents, recursively.

    x:    a string representing the full path to a file or folder.
    mode: an octal number representing the new mode to set.

    """
    for top, folders, files in os.walk(x):
        for folderpath in [os.path.join(top, f) for f in folders]:
            chmod(folderpath, mode)
        for filepath in [os.path.join(top, f) for f in files]:
            chmod(filepath, mode)

def chown_tree(x, user=None, group=None):
    """Recursively sets the owner of a folder and all its contents.

    x:     a string representing the full path to a folder.
    user:  a string representing the name of the new user.
           Alternately, an int representing the uid of the new user.
    group: a string representing the name of the new group.
           Alternately, an int representing the gid of the new group.

    User and group are optional, but giving neither is a no-op.

    """
    if user is None and group is None:
        return
    for top, folders, files in os.walk(x):
        for folderpath in [os.path.join(top, f) for f in folders]:
            chown(folderpath, user, group)
        for filepath in [os.path.join(top, f) for f in files]:
            chown(filepath, user, group)

def chgrp_tree(x, group):
    """Recursively sets the group of a folder and all its contents.

    x:     a string representing the full path to a folder.
    group: a string representing the name of the new group.
           Alternately, an int representing the gid of the new group.

    """
    chown_tree(x, None, group)

def create_folder(name=None, dir=None, method='sha-1'):
    """Creates a folder.
    
    name:   a string representing the name of the folder to create.
    dir:    the name of the folder in which to create the new folder.
    method: a string representing the randon method to use when
            creating a random folder name, if none is given.  'sha-1'
            by default.

    """
    os.mkdir(_get_path(name, dir, method=method))

def rmtree(f):
    """Deletes a folder and all its contents."""
    if not os.path.isdir(f):
        raise OSError, "[%s] is not a folder"
    for path, folders, files in os.walk(f, topdown=False):
        fullfilepaths = [os.path.join(path, x) for x in files]
        fullfolderpaths = [os.path.join(path, x) for x in folders]
        [os.unlink(x) for x in fullfilepaths]
        [os.rmdir(x) for x in fullfolderpaths]
    os.rmdir(f)

def mktree(p):
    """Creates a folder, creating its top-level folders if necessary."""
    p_list = p.split(os.path.sep)
    current = p_list[0]
    p_list = p_list[1:]
    for x in range(len(p_list)):
        if os.path.exists(current):
            if not os.path.isdir(current):
                raise OSError, "[%s] exists but is not a folder" % (current)
            continue
        os.mkdir(current)
        current = os.path.join(current, p_list[x])

delete_folder = rmtree

def create_file(f, overwrite=False):
    """Creates a file.
    
    f:         a string representing the full path to a file to create.
    overwrite: a boolean that, if True, will overwrite an existing
               with a new, empty file.  False by default.

    """
    File(f, overwrite=overwrite).create()

def read_file(fn, size=-1, lock=False):
    """Returns the contents of a file.

    fn:   a string representing the full path to a file to read.
    size: an int representing how many bytes to read from the file.
          If this is omitted or negative, returns all bytes.
    lock: a boolean that, if given, locks a file before reading it,
          unlocking it when finished.

    """
    if lock:
        f = LockFile(fn)
    else:
        f = File(fn)
    return f.read(size)

def write_file(data, fn, overwrite=False, lock=False):
    """Writes data to a file.

    data:      a string representing data to write to a file.
    fn:        a string representing the full path to a file to write
               to.
    overwrite: if True, will overwrite the file.  False by default.
    lock:      if True, locks the file before writing to it.

    If it does not exist, 'f' will be created.  If it does exist, its
    contents will be overwritten.
    
    Normally this operation can fail and leave a new, empty file.  If
    this does occur, write_file will attempt to remove the file if it
    did not previously exist before raising the Exception.  This
    operation will fail silently, if it fails at all.

    """
    if lock:
        f = LockFile(fn, overwrite=overwrite)
    else:
        f = File(fn, overwrite=overwrite)
    f.write(data)

def update_file(data, fn, seek_to=0, overwrite=False, lock=False):
    """Updates a file's contents.

    data:      a string representing data to write to a file.
    fn:        a string representing the full path to a file to update.
    seek_to:   an int representing the position to seek to before
               updating.  Defaults to the beginning of the file.
    overwrite: if True, will overwrite the file.  False by default.
    lock:      if True, locks the file before writing to it.

    """
    if lock:
        f = LockFile(fn, overwrite=overwrite)
    else:
        f = File(fn, overwrite=overwrite)
    if seek_to:
        f.seek(seek_to)
    f.update(data)

def append_file(data, fn, lock=False):
    """Appends data to file.

    data: a string representing data to append to a file.
    fn:   a string representing the full path to a file to append to.
    lock: a boolean that, if given, locks a file before writing to it,
          unlocking it when finished.

    Normally this operation can fail and leave a new, empty file.  If
    this does occur, append_file will attempt to remove the file if it
    did not previously exist before raising the Exception.  This
    operation will fail silently, if it fails at all.

    """
    if lock:
        f = LockFile(fn, overwrite=True)

    else:
        f = File(fn, overwrite=True)
    f.append(data)

def delete_file(f):
    """Deletes a file."""
    File(f).delete()

def ls(folder):
    """Lists folder contents.

    folder: a string representing a folder whose contents will be
            listed.

    os.listdir only returns basenames of files.  'ls' returns a list of
    strings representing full paths to files contained in folder.  Of
    course, this is only possible if 'folder' is an absolute path
    itself.

    """
    rel_paths = [os.path.join(folder, f) for f in os.listdir(folder)]
    return [os.path.abspath(x) for x in rel_paths]

def get_random_file(mode='w+b', method='mkstemp', dir=tempfile.gettempdir(),
                    prefix=tempfile.gettempprefix(), suffix=""):
    """Returns a file object opened on a new temporary file.

    mode:    just like mode in 'open' and 'file'
    method:  a string representing the method by which to generate a
             random filename.  Supported methods include:
               - sha-1
               - mkstemp
             'mkstemp' functions just like tempfile.mkstemp.  'sha-1'
             uses the SHA-1 hash of a "random" string.  'sha-1' is
             better, 'mkstemp' is faster and the default.
    dir:     a string representing the full path to a folder in which
             the temp file will be created.  If not given, the system
             default temporary folder is used.
    prefix:  a string representing a prefix to add to the beginning of
             the base random filename.
    suffix:  a string representing a prefix to add to the end of the
             random filename.

    """
    ###
    # TODO: implement checks for existing files and max # of retries
    ###
    flags = os.O_RDWR | os.O_CREAT | os.O_EXCL
    if hasattr(os, 'O_NOINHERIT'):
        flags |= os.O_NOINHERIT
    if hasattr(os, 'O_NOFOLLOW'):
        flags |= os.O_NOFOLLOW
    if 'b' in mode and hasattr(os, 'O_BINARY'):
        flags |= os.O_BINARY
    p = _get_path(None, dir, prefix=prefix, suffix=suffix, method=method)
    fd = os.open(p, flags, 0600)
    os.close(fd)
    return open(p, mode)

def get_temp_file(mode='w+b', bufsize=-1, persist=True, method='sha-1',
                  dir=tempfile.gettempdir(), prefix=tempfile.gettempprefix(),
                  suffix=""):
    """Returns a file object opened on a new temporary file.

    mode:    just like mode in 'open' and 'file'
    bufsize: just like bufsize in 'open' and 'file'
    persist: an optional boolean that, if given, will keep the
             temporary file on disk after the object is gone.  This is
             the default.
    method:  a string representing the method by which to generate a
             random filename.  Supported methods inclue:
               - sha-1
               - mkstemp
             'mkstemp' functions just like tempfile.mkstemp.  'sha-1'
             uses the SHA-1 hash of a "random" string.  'mkstemp' is
             faster, 'sha-1' is better, and the default.
    dir:     a string representing the full path to a folder in which
             the temp file will be created.  If not given, the system
             default temporary folder is used.
    prefix:  a string representing a prefix to add to the beginning of
             the base random filename.
    suffix:  a string representing a prefix to add to the end of the
             random filename.

    """
    ###
    # It took me a couple hours to really figure out temp files, and I'm still
    # not sure I'm all the way up on it, but here it goes.
    #
    # The first step in creating a temp file/dir is to create a filename that
    # is reasonably unpredictable (can't base it on system variables or state,
    # for example).  Next, you have to open the file under these circumstances:
    #   - must not previously exist (O_CREAT | O_EXCL)
    #   - must not follow symlinks (O_NOFOLLOW)
    #   - must honor/support mode (O_BINARY | O_RDWR)
    #   - must not have any existing data (mode='w+b', which truncates the file)
    #   - must not be accessible by other processes (unlink,
    #                                                O_NOINHERIT | O_TEMPORARY)
    #     * see note below
    #   - mode: 0600
    # * On POSIX, unlinking the file right after it's created keeps other
    #   processes from messing with your file.  On Windows, setting (
    #   (O_NOINHERIT | O_TEMPORARY) does this... though the file is still on
    #   disk.  This also removes the file when it's closed.
    #
    # The TemporaryFile objects in tempfile do all this for you, but sometimes
    # you want randomly-uniquely-named files full of specified data to stick
    # around after you close them or exec new processes.
    #
    # get_temp_file gives you this option with the 'persist' flag.  If
    # 'persist' is False, you just get a NamedTemporaryFile instance.
    # Otherwise, you get a file object opened on a new, randomly-uniquely-named
    # file.  Thus, the name of this function is kind of a misnomer.  But meh.
    ###
    ###
    # TODO: implement checks for existing files and max # of retries
    ###
    flags = os.O_RDWR | os.O_CREAT | os.O_EXCL
    if hasattr(os, 'O_NOINHERIT'):
        flags |= os.O_NOINHERIT
    if hasattr(os, 'O_NOFOLLOW'):
        flags |= os.O_NOFOLLOW
    if 'b' in mode and hasattr(os, 'O_BINARY'):
        flags |= os.O_BINARY
    p = _get_path(None, dir, prefix=prefix, suffix=suffix, method=method)
    fd = os.open(p, flags, 0600)
    return os.fdopen(fd, mode, bufsize)

def get_temp_dir(method='mkstemp', dir=tempfile.gettempdir(), prefix=tempfile.gettempprefix(),
                 suffix=""):
    """Returns a string representing the full path to a new temp dir.

    method:  a string representing the method by which to generate a
             random filename.  Supported methods inclue:
               - sha-1
               - mkdtemp
             'mkstemp' functions just like tempfile.mkstemp.  'sha-1'
             uses the SHA-1 hash of a "random" string.  'sha-1' is
             better, 'mkstemp' is faster, and the default.
    dir:     a string representing the full path to a folder in which
             the temp file will be created.  If not given, the system
             default temporary folder is used.
    prefix:  a string representing a prefix to add to the beginning of
             the base random filename.  The default is 'tmp'.
    suffix:  a string representing a prefix to add to the end of the
             random filename.  The default is 'tmp'.

    """
    p = _get_path(None, dir, prefix=prefix, suffix=suffix, method=method)
    os.mkdir(p, 0700)
    return p

